# 对数周期幂律模型的拐点预测
import math
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from lppls import lppls


def train_lppl(data):
    time = np.arange(0, len(data))
    price = np.log(data.iloc[:, 0].to_list())
    # create observations array (expected format for LPPLS observations)
    observations = np.array([time, price])
    # set the max number for searches to perform before giving-up, the literature suggests 25
    # MAX_SEARCHES = 25
    # # instantiate a new LPPLS model with the data
    lppls_model = lppls.LPPLS(observations=observations)
    # # fit the model to the data and get back the params
    # tc, m, w, a, b, c, c1, c2, O, D = lppls_model.fit(MAX_SEARCHES)
    # lppls_model.plot_fit()
    res = lppls_model.mp_compute_nested_fits(
        workers=8, 
        window_size=120, 
        smallest_window_size=30, 
        outer_increment=1,
        inner_increment=5, 
        max_searches=25, 
        )
    return res, lppls_model.compute_indicators(res)


def plot_bubble(res_df):
    fig, (ax1, ax2) = plt.subplots(nrows=2, ncols=1, sharex=True, figsize=(18, 10))
    ord = res_df['time'].astype('int32')
    ts = [pd.Timestamp.fromordinal(d) for d in ord]
    # plot pos bubbles
    ax1_0 = ax1.twinx()
    ax1.plot(ts, res_df['price'], color='black', linewidth=0.75)
    ax1_0.plot(ts, res_df['pos_conf'], label='bubble indicator (pos)', color='red', alpha=0.5)
    # plot neg bubbles
    ax2_0 = ax2.twinx()
    ax2.plot(ts, res_df['price'], color='black', linewidth=0.75)
    ax2_0.plot(ts, res_df['neg_conf'], label='bubble indicator (neg)', color='green', alpha=0.5)

    ax1.grid(which='major', axis='both', linestyle='--')
    ax2.grid(which='major', axis='both', linestyle='--')
    # set labels
    ax1.set_ylabel('ln(p)')
    ax2.set_ylabel('ln(p)')

    ax1_0.set_ylabel('bubble indicator (pos)')
    ax2_0.set_ylabel('bubble indicator (neg)')

    ax1_0.legend(loc=2)
    ax2_0.legend(loc=2)

    plt.xticks(rotation=45)
    return fig


if __name__ == "__main__":
    data = pd.read_excel('伦敦市场黄金现货价.xlsx', index_col=0)
    model, res_df = train_lppl(data)
    fig = plot_bubble(res_df)
    bubble_df = res_df
    bubble_df['time'] = data['trade_date'].to_list()[-len(bubble_df):]